 #!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
多线程补充缺失的历史最大风速数据脚本
使用线程池加快处理速度
"""

import sys
import requests
import logging
from pathlib import Path
from datetime import datetime, timedelta, date
from typing import List, Dict, Optional
import time
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
import queue

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.absolute()
sys.path.insert(0, str(project_root))

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('multithread_windspeed_fix.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# 线程安全的计数器
class ThreadSafeCounter:
    def __init__(self):
        self._value = 0
        self._lock = threading.Lock()
    
    def increment(self):
        with self._lock:
            self._value += 1
            return self._value
    
    def get_value(self):
        with self._lock:
            return self._value

def get_visual_crossing_api_key():
    """获取Visual Crossing API密钥"""
    try:
        from backend.config.weather_config import WeatherConfig
        api_key = WeatherConfig.get_api_key()
        if api_key and api_key != 'YOUR_VISUAL_CROSSING_API_KEY':
            return api_key
        else:
            logger.error("Visual Crossing API密钥未正确配置")
            return None
    except ImportError:
        logger.error("无法导入WeatherConfig")
        return None

def get_windspeed_from_api(date_str: str, api_key: str) -> Optional[float]:
    """从Visual Crossing API获取指定日期的风速数据"""
    try:
        url = "https://weather.visualcrossing.com/VisualCrossingWebServices/rest/services/timeline/Beijing,China"
        params = {
            'unitGroup': 'metric',
            'include': 'days',
            'key': api_key,
            'contentType': 'json',
            'startDateTime': f"{date_str}T00:00:00",
            'endDateTime': f"{date_str}T23:59:59"
        }
        
        response = requests.get(url, params=params, timeout=10)
        
        if response.status_code == 200:
            data = response.json()
            if 'days' in data and len(data['days']) > 0:
                day_data = data['days'][0]
                windspeed = day_data.get('windspeed')
                if windspeed is not None:
                    logger.info(f"成功获取 {date_str} 的风速数据: {windspeed} km/h")
                    return float(windspeed)
                else:
                    logger.warning(f"{date_str} 的风速数据为空")
                    return None
            else:
                logger.warning(f"{date_str} 没有返回数据")
                return None
        else:
            logger.error(f"API请求失败: {response.status_code} - {response.text}")
            return None
            
    except Exception as e:
        logger.error(f"获取 {date_str} 风速数据时出错: {str(e)}")
        return None

def update_windspeed_record(record_id: int, date_str: str, windspeed: float):
    """更新单条记录的风速数据"""
    try:
        from backend.config.database import get_db_session
        from backend.entities.weather_daily import WeatherDaily
        
        db = get_db_session()
        try:
            # 查找并更新记录
            record = db.query(WeatherDaily).filter(WeatherDaily.id == record_id).first()
            if record:
                record.ws_max = windspeed
                record.updated_at = datetime.now()
                db.commit()
                logger.info(f"更新 {date_str} 的风速数据: {windspeed} km/h")
                return True
            else:
                logger.warning(f"未找到ID为 {record_id} 的记录")
                return False
        finally:
            db.close()
    except Exception as e:
        logger.error(f"更新记录 {record_id} 时出错: {str(e)}")
        return False

def process_record_batch(records_batch: List[tuple], api_key: str, counter: ThreadSafeCounter):
    """处理一批记录"""
    results = []
    
    for record_id, date_str in records_batch:
        try:
            # 获取风速数据
            windspeed = get_windspeed_from_api(date_str, api_key)
            
            if windspeed is not None:
                # 更新数据库
                success = update_windspeed_record(record_id, date_str, windspeed)
                if success:
                    results.append(True)
                else:
                    results.append(False)
            else:
                results.append(False)
            
            # 更新计数器
            current_count = counter.increment()
            if current_count % 50 == 0:
                logger.info(f"进度: {current_count} 条记录已处理")
            
            # 添加延迟避免API限制
            time.sleep(0.5)
            
        except Exception as e:
            logger.error(f"处理记录 {record_id} 时出错: {str(e)}")
            results.append(False)
    
    return results

def get_missing_windspeed_records_batch(batch_size=100):
    """分批获取缺失风速数据的记录"""
    try:
        from backend.config.database import get_db_session
        from backend.entities.weather_daily import WeatherDaily
        from sqlalchemy import and_
        
        db = get_db_session()
        try:
            # 获取缺失风速数据的记录
            missing_records = db.query(WeatherDaily).filter(
                and_(
                    WeatherDaily.t_max.isnot(None),  # 有气象数据
                    WeatherDaily.ws_max.is_(None)    # 但缺失风速
                )
            ).order_by(WeatherDaily.dt).all()
            
            # 转换为 (id, date_str) 格式
            records_data = []
            for record in missing_records:
                date_str = record.dt.strftime('%Y-%m-%d')
                records_data.append((record.id, date_str))
            
            return records_data
            
        finally:
            db.close()
    except Exception as e:
        logger.error(f"获取缺失风速数据记录时出错: {str(e)}")
        return []

def main():
    """主函数"""
    logger.info("开始多线程风速数据补充...")
    
    # 获取API密钥
    api_key = get_visual_crossing_api_key()
    if not api_key:
        logger.error("无法获取API密钥，退出")
        return
    
    # 获取缺失风速数据的记录
    logger.info("获取缺失风速数据的记录...")
    missing_records = get_missing_windspeed_records_batch()
    
    if not missing_records:
        logger.info("没有缺失风速数据的记录")
        return
    
    logger.info(f"找到 {len(missing_records)} 条缺失风速数据的记录")
    
    # 创建线程安全的计数器
    counter = ThreadSafeCounter()
    
    # 将记录分批
    batch_size = 20  # 每个线程处理20条记录
    batches = []
    for i in range(0, len(missing_records), batch_size):
        batch = missing_records[i:i + batch_size]
        batches.append(batch)
    
    logger.info(f"将记录分为 {len(batches)} 个批次，每批 {batch_size} 条记录")
    
    # 使用线程池处理
    max_workers = 5  # 最多5个线程
    total_updated = 0
    
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        # 提交所有批次的任务
        future_to_batch = {
            executor.submit(process_record_batch, batch, api_key, counter): batch 
            for batch in batches
        }
        
        # 处理完成的任务
        for future in as_completed(future_to_batch):
            batch = future_to_batch[future]
            try:
                results = future.result()
                batch_updated = sum(results)
                total_updated += batch_updated
                logger.info(f"批次完成: {batch_updated}/{len(batch)} 条记录更新成功")
            except Exception as e:
                logger.error(f"批次处理失败: {str(e)}")
    
    logger.info(f"多线程风速数据补充完成！总共更新了 {total_updated} 条记录")

if __name__ == "__main__":
    main()